import win32com.client
import pythoncom
import datetime


class TaskManager:
    def __init__(self):
        self.task_name = "AutoShutdownTask"

    def create_scheduled_task(self, shutdown_datetime, schedule_type, weekdays):
        """创建Windows计划任务"""
        pythoncom.CoInitialize()
        try:
            scheduler = win32com.client.Dispatch('Schedule.Service')
            scheduler.Connect()
            root_folder = scheduler.GetFolder('\\')

            # 创建任务定义
            task_def = scheduler.NewTask(0)
            task_def.RegistrationInfo.Description = "自动关机任务"
            task_def.Settings.Enabled = True
            task_def.Settings.Hidden = False

            # 根据计划类型创建对应触发器
            if schedule_type == "单次":
                # 单次触发器
                trigger = task_def.Triggers.Create(1)  # 1 = 单次触发器
                trigger.StartBoundary = shutdown_datetime.strftime("%Y-%m-%dT%H:%M:%S")
            elif schedule_type == "每天":
                # 每天触发器
                trigger = task_def.Triggers.Create(2)  # 2 = 每天触发器
                trigger.StartBoundary = shutdown_datetime.strftime("%Y-%m-%dT%H:%M:%S")
                trigger.DaysInterval = 1  # 每天执行一次
            else:  # 每周
                # 每周触发器
                trigger = task_def.Triggers.Create(3)  # 3 = 每周触发器
                trigger.StartBoundary = shutdown_datetime.strftime("%Y-%m-%dT%H:%M:%S")

                # 计算星期掩码（周一=2, 周二=4, 周三=8, 周四=16, 周五=32, 周六=64, 周日=1）
                days_mask = 0
                days_map = {"周一": 2, "周二": 4, "周三": 8, "周四": 16, "周五": 32, "周六": 64, "周日": 1}
                for day, enabled in weekdays.items():
                    if enabled:
                        days_mask |= days_map[day]

                trigger.DaysOfWeek = days_mask
                trigger.WeeksInterval = 1  # 每周重复执行

            # 添加关机操作
            action = task_def.Actions.Create(0)  # 0 = 执行操作
            action.Path = "shutdown"
            action.Arguments = "/s /t 0"

            # 注册任务（覆盖已有任务）
            root_folder.RegisterTaskDefinition(
                self.task_name,
                task_def,
                6,  # TASK_CREATE_OR_UPDATE
                "",  # 用户名（空表示使用当前用户）
                "",  # 密码（空表示使用当前用户）
                3,  # TASK_LOGON_SERVICE_ACCOUNT（以系统服务身份运行）
                ""  # 安全描述符（空表示使用默认）
            )
        finally:
            pythoncom.CoUninitialize()

    def delete_scheduled_task(self):
        """删除Windows计划任务"""
        pythoncom.CoInitialize()
        try:
            scheduler = win32com.client.Dispatch('Schedule.Service')
            scheduler.Connect()
            root_folder = scheduler.GetFolder('\\')
            root_folder.DeleteTask(self.task_name, 0)
        except Exception:
            # 如果任务不存在，忽略错误
            pass
        finally:
            pythoncom.CoUninitialize()

    def check_existing_task(self):
        """检查是否存在已创建的计划任务"""
        pythoncom.CoInitialize()
        try:
            scheduler = win32com.client.Dispatch('Schedule.Service')
            scheduler.Connect()
            root_folder = scheduler.GetFolder('\\')
            # 尝试获取任务，存在则返回True
            root_folder.GetTask(self.task_name)
            return True
        except Exception:
            return False
        finally:
            pythoncom.CoUninitialize()